#pragma once

#include <vector>
#include <memory>
#include "llvm/IR/Value.h"
#include "type.h"
#include "lexer.h"

enum class OpCode
{
    add,
    sub,
    mul,
    div,
    equal_equal,
    not_equal,
    less,
    less_equal,
    greater,
    greater_equal
};

// 进行声明
class Program;
class Expr;
class DeclStmt;
class BlockStmt;
class ForStmt;
class BreakStmt;
class ContinueStmt;
class VariableDecl;
class IfStmt;
class VariableAccessExpr;
class BinaryExpr;
class AssignExpr;
class NumberExpr;

// 访问者模式
class Visitor
{
public:
    virtual ~Visitor() {}
    virtual llvm::Value *VisitProgram(Program *p) = 0;
    virtual llvm::Value *VisitDeclStmt(DeclStmt *decl) = 0;
    virtual llvm::Value *VisitBlockStmt(BlockStmt *block) = 0;
    virtual llvm::Value *VisitIfStmt(IfStmt *ifStmt) = 0;
    virtual llvm::Value *VisitForStmt(ForStmt *forStmt) = 0;
    virtual llvm::Value *VisitBreakStmt(BreakStmt *breakStmt) = 0;
    virtual llvm::Value *VisitContinueStmt(ContinueStmt *continueStmt) = 0;
    virtual llvm::Value *VisitVariableDecl(VariableDecl *decl) = 0;
    virtual llvm::Value *VisitVariableAccessExpr(VariableAccessExpr *varaccExpr) = 0;
    virtual llvm::Value *VisitAssignExpr(AssignExpr *assignExpr) = 0;
    virtual llvm::Value *VisitBinaryExpr(BinaryExpr *binaryExpr) = 0;
    virtual llvm::Value *VisitNumberExpr(NumberExpr *factorExpr) = 0;
};

// 语法树的公共节点
class ASTNode
{
public:
    enum Kind
    {
        ND_BlockStmt,
        ND_ForStmt,
        ND_IfStmt,
        ND_BreakStmt,
        ND_ContinueStmt,
        ND_DeclStmt,
        ND_VariableDecl,
        ND_BinaryExpr,
        ND_NumberExpr,
        ND_VariableAccessExpr,
        ND_AssignExpr
    };

private:
    const Kind kind;

public:
    ASTNode(Kind kind) : kind(kind) {}
    virtual ~ASTNode() {}
    virtual llvm::Value *Accept(Visitor *v) { return nullptr; } // 通过虚函数的特性完成分发的功能
    const Kind getKind() const { return kind; }

public:
    CType *type;
    Token token;
};

// 声明语句节点
class DeclStmt : public ASTNode
{
public:
    DeclStmt() : ASTNode(ND_DeclStmt) {}

    llvm::Value *Accept(Visitor *v) override
    {
        return v->VisitDeclStmt(this);
    }

    static bool classof(const ASTNode *node) // 判断是否能进行指针的强转
    {
        return node->getKind() == ND_DeclStmt;
    }

public:
    std::vector<std::shared_ptr<ASTNode>> nodeVec;
};

// 块语句节点
class BlockStmt : public ASTNode
{
public:
    BlockStmt() : ASTNode(ND_BlockStmt) {}

    llvm::Value *Accept(Visitor *v) override
    {
        return v->VisitBlockStmt(this);
    }

    static bool classof(const ASTNode *node) // 判断是否能进行指针的强转
    {
        return node->getKind() == ND_BlockStmt;
    }

public:
    std::vector<std::shared_ptr<ASTNode>> nodeVec;
};

// 条件判断节点
class IfStmt : public ASTNode
{
public:
    IfStmt() : ASTNode(ND_IfStmt) {}

    llvm::Value *Accept(Visitor *v) override
    {
        return v->VisitIfStmt(this);
    }

    static bool classof(const ASTNode *node) // 判断是否能进行指针的强转
    {
        return node->getKind() == ND_IfStmt;
    }

public:
    std::shared_ptr<ASTNode> condNode;
    std::shared_ptr<ASTNode> thenNode;
    std::shared_ptr<ASTNode> elseNode;
};

// 循环节点
class ForStmt : public ASTNode
{
public:
    ForStmt() : ASTNode(ND_ForStmt) {}

    llvm::Value *Accept(Visitor *v) override
    {
        return v->VisitForStmt(this);
    }

    static bool classof(const ASTNode *node) // 判断是否能进行指针的强转
    {
        return node->getKind() == ND_ForStmt;
    }

public:
    std::shared_ptr<ASTNode> initNode;
    std::shared_ptr<ASTNode> condNode;
    std::shared_ptr<ASTNode> incNode;
    std::shared_ptr<ASTNode> bodyNode;
};

class BreakStmt : public ASTNode
{
public:
    BreakStmt() : ASTNode(ND_BreakStmt) {}

    llvm::Value *Accept(Visitor *v) override
    {
        return v->VisitBreakStmt(this);
    }

    static bool classof(const ASTNode *node) // 判断是否能进行指针的强转
    {
        return node->getKind() == ND_BreakStmt;
    }

public:
    std::shared_ptr<ASTNode> target;
};

class ContinueStmt : public ASTNode
{
public:
    ContinueStmt() : ASTNode(ND_ContinueStmt) {}

    llvm::Value *Accept(Visitor *v) override
    {
        return v->VisitContinueStmt(this);
    }

    static bool classof(const ASTNode *node) // 判断是否能进行指针的强转
    {
        return node->getKind() == ND_ContinueStmt;
    }

public:
    std::shared_ptr<ASTNode> target;
};


// 变量声明节点
class VariableDecl : public ASTNode
{
public:
    VariableDecl() : ASTNode(ND_VariableDecl) {}

    llvm::Value *Accept(Visitor *v) override
    {
        return v->VisitVariableDecl(this);
    }

    static bool classof(const ASTNode *node) // 判断是否能进行指针的强转
    {
        return node->getKind() == ND_VariableDecl;
    }
};

// 二元表达式节点
class BinaryExpr : public ASTNode
{
public:
    BinaryExpr() : ASTNode(ND_BinaryExpr) {}

    llvm::Value *Accept(Visitor *v) override
    {
        return v->VisitBinaryExpr(this);
    }

    static bool classof(const ASTNode *node) // 判断是否能进行指针的强转
    {
        return node->getKind() == ND_BinaryExpr;
    }

public:
    OpCode op;
    std::shared_ptr<ASTNode> left;
    std::shared_ptr<ASTNode> right;
};

// 赋值表达式节点
class AssignExpr : public ASTNode
{
public:
    AssignExpr() : ASTNode(ND_AssignExpr) {}

    llvm::Value *Accept(Visitor *v) override
    {
        return v->VisitAssignExpr(this);
    }

    static bool classof(const ASTNode *node) // 判断是否能进行指针的强转
    {
        return node->getKind() == ND_AssignExpr;
    }

public:
    std::shared_ptr<ASTNode> left;
    std::shared_ptr<ASTNode> right;
};

// 数字表达式节点
class NumberExpr : public ASTNode
{
public:
    NumberExpr() : ASTNode(ND_NumberExpr) {}

    llvm::Value *Accept(Visitor *v) override
    {
        return v->VisitNumberExpr(this);
    }

    static bool classof(const ASTNode *node) // 判断是否能进行指针的强转
    {
        return node->getKind() == ND_NumberExpr;
    }
};

// 变量访问节点
class VariableAccessExpr : public ASTNode // 变量表达式
{
public:
    VariableAccessExpr() : ASTNode(ND_VariableAccessExpr) {}

    llvm::Value *Accept(Visitor *v) override
    {
        return v->VisitVariableAccessExpr(this);
    }

    static bool classof(const ASTNode *node) // 判断是否能进行指针的强转
    {
        return node->getKind() == ND_VariableAccessExpr;
    }
};

// 目标程序
class Program
{
public:
    std::vector<std::shared_ptr<ASTNode>> nodeVec;
};
